import numpy as np
import random
from Translate import Translate
from Evaluate import Evaluate

class KMEANS(object):

    def __init__(self, k):
        self.k = k
        self.labels = np.array([0.0]*self.k)

    def fit(self, traindata, trainlabel):
        n = len(traindata)
        r = random.sample(range(n), self.k)
        mu = traindata[r]
        dis = np.zeros((n,self.k))
        for iter_num in range(1000):
            for i in range(self.k):
                dis[:,i] = np.sum((traindata - mu[i]) ** 2, 1)
            new_mu = np.zeros((self.k, len(traindata[0])))
            num = [0] * self.k
            for i in range(n):
                tmp = np.where(dis[i] == np.min(dis[i]))[0][0]
                new_mu[tmp] += traindata[i]
                num[tmp] += 1
            # print num
            for i in range(self.k):
                new_mu[i] /= num[i]
            if np.max(abs(new_mu - mu)) < 1e-6:
                break
            # print iter_num
            # print mu
            mu = new_mu
        self.mu = mu

        dis = np.zeros((len(traindata), self.k))
        for i in range(self.k):
            dis[:,i] = np.sum((traindata-self.mu[i])**2, 1)
        for i in range(self.k):
            tmp = [trainlabel[x] for x in range(len(traindata)) if np.where(dis[x] == np.min(dis[x]))[0][0] == i]
            self.labels[i] = np.mean(tmp)

    def predict(self, data):
        dis = np.zeros((len(data), self.k))
        for i in range(self.k):
            dis[:,i] = np.sum((data-self.mu[i])**2, 1)
        predict = np.array([self.labels[np.where(dis[x] == np.min(dis[x]))[0][0]] for x in range(len(data))])
        # print predict
        return predict

def train_kmeans(isEvaluate = True, task = 2):
    evaluate = Evaluate()

    random.seed(26)

    if task == 0 or task == 1:
        data0 = np.load('./data_task1/traindata_part0_normalizing.npz')
        data1 = np.load('./data_task1/traindata_part1_normalizing.npz')
        if not isEvaluate:
            traindata0 = data0['datas']
            traindata1 = data1['datas']
            trainlabel0 = data0['labels']
            trainlabel1 = data1['labels']
            predictlabel0 = np.zeros((6, len(traindata0), 6))
            predictlabel1 = np.zeros((6, len(traindata1), 6))
        else:
            traindata0 = data0['datas'][0:-7]
            traindata1 = data1['datas'][0:-7]
            trainlabel0 = data0['labels'][:,0:-7,:]
            trainlabel1 = data1['labels'][:,0:-7,:]
            predictlabel0 = np.zeros((6, len(traindata0), 6))
            predictlabel1 = np.zeros((6, len(traindata1), 6))

        models = {}

        for i in range(6):
            for j in range(6):
                model = KMEANS(k=6)
                Y = trainlabel0[i,:,j]
                model.fit(traindata0, Y)
                models[(i, j, 0)] = model
                predictlabel0[i,:,j] = model.predict(traindata0) * 0.91

                model = KMEANS(k=6)
                Y = trainlabel1[i,:,j]
                model.fit(traindata1, Y)
                models[(i, j, 1)] = model
                predictlabel1[i,:,j] = model.predict(traindata1) * 0.93

        MAPE_train_0 = evaluate.evaluate_test(predictlabel0, trainlabel0)
        MAPE_train_1 = evaluate.evaluate_test(predictlabel1, trainlabel1)
        print 'travel time MAPE_train: ' + str(MAPE_train_0) + ' ' + str(MAPE_train_1)


        if not isEvaluate:
            data0 = np.load('./data_task1/testdata_part0_normalizing.npz')
            data1 = np.load('./data_task1/testdata_part1_normalizing.npz')
            testdata0 = data0['tests']
            testdata1 = data1['tests']
            testlabel = np.zeros((6, 7, 12))
        else:
            testdata0 = data0['datas'][-7:]
            testdata1 = data1['datas'][-7:]
            testlabel = np.zeros((6, 7, 12))
            reslabel = np.concatenate((data0['labels'][:,-7:,:], data1['labels'][:,-7:,:]), axis=2)

        for i in range(6):
            for j in range(6):
                testlabel[i,:,j] = models[(i, j, 0)].predict(testdata0) * 0.91
                testlabel[i,:,j+6] = models[(i, j, 1)].predict(testdata1) * 0.93

        # print testlabel

        if not isEvaluate:
            trans = Translate()
            trans.translate(testlabel, path = './result_svm', task = 1)
        else:
            MAPE = evaluate.evaluate_test(testlabel, reslabel)
            print 'travel time MAPE: ' + str(MAPE)

        if task == 1:
            return testlabel


    if task == 0 or task == 2:
        data0 = np.load('./data_task2/traindata_part0_nominmaxscaler.npz')
        data1 = np.load('./data_task2/traindata_part1_nominmaxscaler.npz')
        if not isEvaluate:
            traindata0 = data0['datas']
            traindata1 = data1['datas']
            trainlabel0 = data0['labels']
            trainlabel1 = data1['labels']
            predictlabel0 = np.zeros((5, len(traindata0), 6))
            predictlabel1 = np.zeros((5, len(traindata1), 6))
        else:
            traindata0 = data0['datas'][0:-7]
            traindata1 = data1['datas'][0:-7]
            trainlabel0 = data0['labels'][:,0:-7,:]
            trainlabel1 = data1['labels'][:,0:-7,:]
            predictlabel0 = np.zeros((5, len(traindata0), 6))
            predictlabel1 = np.zeros((5, len(traindata1), 6))

        models = {}

        for i in range(5):
            for j in range(6):
                model = KMEANS(k=5)
                Y = trainlabel0[i,:,j]
                model.fit(traindata0, Y)
                models[(i, j, 0)] = model
                predictlabel0[i,:,j] = model.predict(traindata0) * 0.93

                model = KMEANS(k=5)
                Y = trainlabel1[i,:,j]
                model.fit(traindata1, Y)
                models[(i, j, 1)] = model
                predictlabel1[i,:,j] = model.predict(traindata1) * 0.87

        MAPE_train_0 = evaluate.evaluate_test(predictlabel0, trainlabel0)
        MAPE_train_1 = evaluate.evaluate_test(predictlabel1, trainlabel1)
        print 'traffic volume MAPE_train_0/1: ' + str(MAPE_train_0) + ' ' + str(MAPE_train_1)
        
        if not isEvaluate:
            data0 = np.load('./data_task2/testdata_part0_nominmaxscaler.npz')
            data1 = np.load('./data_task2/testdata_part1_nominmaxscaler.npz')
            testdata0 = data0['tests']
            testdata1 = data1['tests']
            testlabel = np.zeros((5, 7, 12))
        else:
            testdata0 = data0['datas'][-7:]
            testdata1 = data1['datas'][-7:]
            testlabel = np.zeros((5, 7, 12))
            reslabel = np.concatenate((data0['labels'][:,-7:,:], data1['labels'][:,-7:,:]), axis=2)

        for i in range(5):
            for j in range(6):
                testlabel[i,:,j] = models[(i, j, 0)].predict(testdata0) * 0.94
                testlabel[i,:,j+6] = models[(i, j, 1)].predict(testdata1) * 0.88

        # print testlabel

        if not isEvaluate:
            trans = Translate()
            trans.translate(testlabel, path = './result_kmeans', task = 2)
        else:
            MAPE = evaluate.evaluate_test(testlabel, reslabel)
            print 'traffic volume MAPE: ' + str(MAPE)

        if task == 2:
            return testlabel

if __name__ == '__main__':
    train_kmeans(isEvaluate = True, task = 2)
